import numpy as np
import matplotlib.pyplot as plt
import random
from scipy.interpolate import CubicSpline
from scipy.interpolate import BSpline
from scipy.ndimage import gaussian_filter1d

class APF_RRTStar:
    def __init__(self, start, goal, obstacle_list, x_limits, y_limits, robot_radius=0.5):
        self.start = start
        self.goal = goal
        self.obstacle_list = obstacle_list
        self.x_limits = x_limits
        self.y_limits = y_limits
        self.robot_radius = robot_radius  # Added robot radius
        self.tree_start = [start]
        self.tree_goal = [goal]
        self.parent_start = {start: None}
        self.parent_goal = {goal: None}
        self.goal_reached = False
        self.step_size = 0.8
        self.max_iterations = 10000  # Reasonable iteration limit
        self.connection_point = None  # To store the connection point of both trees

    def distance(self, point1, point2):
        return np.linalg.norm(np.array(point1) - np.array(point2))

    def is_collision_free(self, point):
        for (ox, oy, radius) in self.obstacle_list:
            if self.distance((ox, oy), point) <= (radius + self.robot_radius):
                return False
        return True

    def get_random_point(self, goal_bias):
        if random.random() < goal_bias:
            return self.goal
        while True:
            point = (random.uniform(self.x_limits[0], self.x_limits[1]),
                     random.uniform(self.y_limits[0], self.y_limits[1]))
            if self.is_collision_free(point):
                return point

    def get_nearest_node(self, point, tree):
        distances = [self.distance(node, point) for node in tree]
        nearest_index = np.argmin(distances)
        return tree[nearest_index]

    def generate_new_node(self, from_node, to_point):
        direction = np.array(to_point) - np.array(from_node)
        length = np.linalg.norm(direction)
        if length == 0:
            return None
        direction = direction / length
        new_node = tuple(np.array(from_node) + self.step_size * direction)
        if self.is_collision_free(new_node):
            return new_node
        return None

    def connect_trees(self, node_from_start, node_from_goal):
        return self.distance(node_from_start, node_from_goal) < self.step_size

    def potential_field(self, point):
        attractive_force = np.array(self.goal) - np.array(point)
        repulsive_force = np.zeros(2)
        for (ox, oy, radius) in self.obstacle_list:
            obstacle_vector = np.array(point) - np.array((ox, oy))
            distance = np.linalg.norm(obstacle_vector)
            if distance <= (radius + self.robot_radius):
                repulsive_force += obstacle_vector / (distance ** 2)
        return attractive_force + repulsive_force

    def search_path(self, goal_bias=0.1):
        iterations = 0
        while not self.goal_reached and iterations < self.max_iterations:
            iterations += 1

            # Extend the start tree
            random_point_start = self.get_random_point(goal_bias)
            nearest_node_start = self.get_nearest_node(random_point_start, self.tree_start)
            new_node_start = self.generate_new_node(nearest_node_start, random_point_start)

            if new_node_start and self.is_collision_free(new_node_start):
                self.tree_start.append(new_node_start)
                self.parent_start[new_node_start] = nearest_node_start

                nearest_node_goal = self.get_nearest_node(new_node_start, self.tree_goal)
                if self.connect_trees(new_node_start, nearest_node_goal):
                    self.goal_reached = True
                    self.connection_point = (new_node_start, nearest_node_goal)

            # Extend the goal tree
            random_point_goal = self.get_random_point(goal_bias)
            nearest_node_goal = self.get_nearest_node(random_point_goal, self.tree_goal)
            new_node_goal = self.generate_new_node(nearest_node_goal, random_point_goal)

            if new_node_goal and self.is_collision_free(new_node_goal):
                self.tree_goal.append(new_node_goal)
                self.parent_goal[new_node_goal] = nearest_node_goal

                nearest_node_start = self.get_nearest_node(new_node_goal, self.tree_start)
                if self.connect_trees(new_node_goal, nearest_node_start):
                    self.goal_reached = True
                    self.connection_point = (nearest_node_start, new_node_goal)

        if not self.goal_reached:
            print("Failed to find a path within the iteration limit.")
            return []

        # Construct the path
        path_start = []
        path_goal = []
        node = self.connection_point[0]
        while node is not None:
            path_start.append(node)
            node = self.parent_start[node]

        node = self.connection_point[1]
        while node is not None:
            path_goal.append(node)
            node = self.parent_goal[node]

        path_start.reverse()
        path = [self.start]+path_start + path_goal+[self.goal]
        return path

    def smooth_path(self, path):
        if len(path) < 3:  # No need to smooth if path has less than 3 points
            return path

        x = [point[0] for point in path]
        y = [point[1] for point in path]

        # 调整采样点数量
        num_points = len(path) * 5  # 增加采样点数量
        t = np.linspace(0, len(path) - 1, len(path))
        t_new = np.linspace(0, len(path) - 1, num=num_points)

        # 使用三次样条进行插值
        cs_x = CubicSpline(t, x)
        cs_y = CubicSpline(t, y)
        smooth_x = cs_x(t_new)
        smooth_y = cs_y(t_new)

        # 使用高斯滤波器平滑路径
        smooth_x = gaussian_filter1d(smooth_x, sigma=10)
        smooth_y = gaussian_filter1d(smooth_y, sigma=10)

        smooth_path = list(zip(smooth_x, smooth_y))

        return smooth_path


    def draw_graph(self, path=None, smooth_path=None):
        plt.figure()
        for (ox, oy, radius) in self.obstacle_list:
            circle = plt.Circle((ox, oy), radius, color='r', alpha=0.5)
            plt.gca().add_patch(circle)
        for node in self.tree_start:
            if self.parent_start[node] is not None:
                plt.plot([node[0], self.parent_start[node][0]], [node[1], self.parent_start[node][1]], "g-")
        for node in self.tree_goal:
            if self.parent_goal[node] is not None:
                plt.plot([node[0], self.parent_goal[node][0]], [node[1], self.parent_goal[node][1]], "b-")
        if path:
            plt.plot([x for (x, y) in path], [y for (x, y) in path], '-o', color='orange', label='Original Path')
        if smooth_path:
            plt.plot([x for (x, y) in smooth_path], [y for (x, y) in smooth_path], '-', color='purple',
                     label='Smoothed Path')
        plt.plot(self.start[0], self.start[1], "bo")
        plt.plot(self.goal[0], self.goal[1], "ro")
        plt.xlim(self.x_limits)
        plt.ylim(self.y_limits)
        plt.grid(True)
        plt.legend()
        # plt.show()


def main():
    start = (0, 0)
    goal = (10, 10)
    obstacle_list = [
        (5, 5, 0.2),
        (7, 6, 1),
        (2, 3, 1.2),
        (4, 6, 0.4),
        (8, 2, 0.5),
        (2, 8, 0.4)
    ]
    x_limits = (0, 10)
    y_limits = (0, 10)
    robot_radius = 0.25  # Set the robot radius

    apf_rrt_star = APF_RRTStar(start, goal, obstacle_list, x_limits, y_limits, robot_radius)
    path = apf_rrt_star.search_path()

    if path:
        smooth_path = apf_rrt_star.smooth_path(path)
        apf_rrt_star.draw_graph(path, smooth_path)
        print(smooth_path)
    else:
        print("No path found.")


class RTT_planer:
    def __init__(self, start, goal, obstacle_list, num_trials=10):
        self.start = start
        self.goal = goal
        self.obstacle_list = obstacle_list
        self.num_trials = num_trials  # 进行多次规划的次数

    def way_plan(self):
        x_limits = (0, 10)
        y_limits = (0, 10)
        robot_radius = 0.5  # Set the robot radius

        best_path = None
        best_smooth_path = None
        best_rrt = None
        min_length = float('inf')  # 初始化为正无穷

        for _ in range(self.num_trials):
            apf_rrt_star = APF_RRTStar(self.start, self.goal, self.obstacle_list, x_limits, y_limits, robot_radius)
            path = apf_rrt_star.search_path()

            if path:
                smooth_path = apf_rrt_star.smooth_path(path)
                # 计算路径的总长度
                length = sum(np.linalg.norm(np.array(smooth_path[i]) - np.array(smooth_path[i+1]))
                             for i in range(len(smooth_path) - 1))

                # 比较并更新最优路径
                if length < min_length:
                    min_length = length
                    best_path = path
                    best_smooth_path = smooth_path
                    best_rrt = apf_rrt_star

        if best_path:
            # apf_rrt_star = APF_RRTStar(self.start, self.goal, self.obstacle_list, x_limits, y_limits, robot_radius)
            best_rrt.draw_graph(best_path, best_smooth_path)
            return best_smooth_path
        else:
            print("No path found.")
            return None

if __name__ == '__main__':
    planner = RTT_planer(start=(0, 0), goal=(10, 10), obstacle_list=[
        (5, 3, 0.6),
        (7, 6, 1),
        (2, 3, 1.2),
        (4, 6, 0.4),
        (8, 2, 0.5),
        (2, 8, 0.4),
        (8,9,0.5)
    ], num_trials=10)  # 进行10次规划
    best_path = planner.way_plan()
    if best_path:
        print("Best path found:", best_path)
        print(len(best_path))
    plt.show()




# if __name__ == '__main__':
#     main()
#     plt.show()
